{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "329550c1",
   "metadata": {},
   "source": [
    "参考: [一起实现一个Baby Triton](https://zhuanlan.zhihu.com/p/709844371)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08e696c3",
   "metadata": {},
   "source": [
    "# 1 example jit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d9f80926",
   "metadata": {},
   "outputs": [],
   "source": [
    "def jit(fn):\n",
    "\tdef wrapper(*args, **kwargs):\n",
    "\t\tprint(\"jit\")\n",
    "\t\treturn fn(*args, **kwargs)\n",
    "\treturn wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f1e679b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit\n",
    "def add():\n",
    "\tprint(\"add\")\n",
    "\n",
    "# add()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6c0c7e31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "now\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.jit.<locals>.inner()>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def jit(fn):\n",
    "\tprint(\"now\")\n",
    "\tdef inner():\n",
    "\t\tprint(\"jit is called\")\n",
    "\treturn inner\n",
    "\n",
    "@jit\n",
    "def add():\n",
    "\tprint(\"add\")\n",
    "\n",
    "add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fa1de939",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2\n"
     ]
    }
   ],
   "source": [
    "print(1, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02c268de",
   "metadata": {},
   "source": [
    "# 2 词法分析与语法分析"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "76cbe140",
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "import ast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "57f124fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "now jit\n",
      "Module(\n",
      "  body=[\n",
      "    FunctionDef(\n",
      "      name='add',\n",
      "      args=arguments(\n",
      "        posonlyargs=[],\n",
      "        args=[],\n",
      "        kwonlyargs=[],\n",
      "        kw_defaults=[],\n",
      "        defaults=[]),\n",
      "      body=[\n",
      "        Expr(\n",
      "          value=Call(\n",
      "            func=Name(id='print', ctx=Load()),\n",
      "            args=[\n",
      "              Constant(value='add')],\n",
      "            keywords=[]))],\n",
      "      decorator_list=[\n",
      "        Name(id='jit', ctx=Load())],\n",
      "      type_params=[])],\n",
      "  type_ignores=[])\n"
     ]
    }
   ],
   "source": [
    "def jit(fn):\n",
    "\tprint(\"now jit\")\n",
    "\tdef inner():\n",
    "\t\tfn_src = inspect.getsource(fn)\n",
    "\t\tfn_ast = ast.parse(fn_src)\n",
    "\t\tprint(ast.dump(fn_ast, indent=2))\n",
    "\treturn inner\n",
    "\n",
    "@jit\n",
    "def add():\n",
    "\tprint(\"add\")\n",
    "\n",
    "add()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "66d6ec6e",
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (<unknown>, line 1)",
     "output_type": "error",
     "traceback": [
      "Traceback \u001b[1;36m(most recent call last)\u001b[0m:\n",
      "\u001b[0m  File \u001b[0;32md:\\Miniconda\\envs\\my_env_torch\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3577\u001b[0m in \u001b[0;35mrun_code\u001b[0m\n    exec(code_obj, self.user_global_ns, self.user_ns)\u001b[0m\n",
      "\u001b[0m  Cell \u001b[0;32mIn[6], line 9\u001b[0m\n    fn_ast = ast.parse(fn_src, mode='eval')\u001b[0m\n",
      "\u001b[1;36m  File \u001b[1;32md:\\Miniconda\\envs\\my_env_torch\\Lib\\ast.py:52\u001b[1;36m in \u001b[1;35mparse\u001b[1;36m\n\u001b[1;33m    return compile(source, filename, mode, flags,\u001b[1;36m\n",
      "\u001b[1;36m  File \u001b[1;32m<unknown>:1\u001b[1;36m\u001b[0m\n\u001b[1;33m    def add():\u001b[0m\n\u001b[1;37m    ^\u001b[0m\n\u001b[1;31mSyntaxError\u001b[0m\u001b[1;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "def add():\n",
    "\tprint(\"add\")\n",
    "\n",
    "# mod = inspect.getmodule(add)\n",
    "# print(mod, type(mod))\n",
    "# inspect.getsource(mod)\n",
    "\n",
    "fn_src = inspect.getsource(add)\n",
    "fn_ast = ast.parse(fn_src, mode='eval')\n",
    "# ast.unparse(fn_ast)\n",
    "# ast.dump(fn_ast, indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "dd3e5cb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sub():\n",
    "\tprint(\"sub\")\n",
    "\n",
    "fn_ast = ast.parse(inspect.getsource(sub))\n",
    "result = ast.dump(fn_ast)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "74d529c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Stevejsl\\.ninetoothed\n"
     ]
    }
   ],
   "source": [
    "import pathlib\n",
    "CACHE_DIR = pathlib.Path.home() / \".ninetoothed\"\n",
    "print(CACHE_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c368088",
   "metadata": {},
   "source": [
    "# 3 跳过语法检查，隔离JIT的实现\n",
    "\n",
    "代码生成的输入是AST，目标是IR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f9bd7471",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Module(body=[FunctionDef(name='add', args=arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]), body=[Expr(value=Call(func=Name(id='print', ctx=Load()), args=[Constant(value='add')], keywords=[]))], decorator_list=[Call(func=Name(id='jit', ctx=Load()), args=[], keywords=[keyword(arg='target', value=Constant(value='cpu'))])], type_params=[])], type_ignores=[])\n"
     ]
    }
   ],
   "source": [
    "import inspect\n",
    "import ast\n",
    "\n",
    "def jit(target=\"cpu\"):\n",
    "\tassert target in [\"cpu\", \"gpu\"]\n",
    "\tdef inner(fn):\n",
    "\t\treturn JIT(fn, target=target)\n",
    "\treturn inner\n",
    "\n",
    "class JIT:\n",
    "\tdef __init__(self, fn, target=\"cpu\"):\n",
    "\t\tself.fn = fn\n",
    "\t\tself.target = target\n",
    "\t\n",
    "\tdef __call__(self, *args, **kwargs):\n",
    "\t\tfn_src = inspect.getsource(self.fn)\n",
    "\t\tfn_ast = ast.parse(fn_src)\n",
    "\t\tprint(ast.dump(fn_ast))\n",
    "\n",
    "@jit(target=\"cpu\")\n",
    "def add():\n",
    "\tprint(\"add\")\n",
    "\n",
    "add()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "165ef0f7",
   "metadata": {},
   "source": [
    "# 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "ecf0ae50",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'tvm'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[29], line 3\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01minspect\u001b[39;00m\n\u001b[0;32m      2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mast\u001b[39;00m\n\u001b[1;32m----> 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtvm\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m relax \u001b[38;5;28;01mas\u001b[39;00m rx\n\u001b[0;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtvm\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mscript\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m relax \u001b[38;5;28;01mas\u001b[39;00m R\n\u001b[0;32m      5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtvm\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mscript\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mir_builder\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m relax \u001b[38;5;28;01mas\u001b[39;00m relax_builder, ir \u001b[38;5;28;01mas\u001b[39;00m I, IRBuilder \u001b[38;5;28;01mas\u001b[39;00m IB\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'tvm'"
     ]
    }
   ],
   "source": [
    "import inspect\n",
    "import ast\n",
    "from tvm import relax as rx\n",
    "from tvm.script import relax as R\n",
    "from tvm.script.ir_builder import relax as relax_builder, ir as I, IRBuilder as IB\n",
    "def jit(target=\"cpu\"):\n",
    "    assert target in [\"cpu\", \"gpu\"]\n",
    "    def inner(fn):\n",
    "        return JIT(fn, target=target)\n",
    "    return inner\n",
    "\n",
    "class JIT:\n",
    "    def __init__(self, fn, target=\"cpu\"):\n",
    "        self.fn = fn\n",
    "        self.target = target\n",
    "    \n",
    "    def __call__(self, *args, **kwargs):\n",
    "        fn_src = inspect.getsource(self.fn)\n",
    "        fn_ast = ast.parse(fn_src)\n",
    "        print(ast.dump(fn_ast))\n",
    "        code_generator = CodeGenerator(fn_ast, self.target)\n",
    "        code_generator.code_gen()\n",
    "\n",
    "class CodeGenerator(ast.NodeVisitor):\n",
    "    def __init__(self, fn_ast, target):\n",
    "        self.fn_ast = fn_ast\n",
    "        self.target = target\n",
    "        self.ib = IB()\n",
    "        self.ir_module = None\n",
    "        self.entry = None\n",
    "        self.ret = None\n",
    "    \n",
    "    def code_gen(self):\n",
    "        with self.ib:\n",
    "            self.visit(self.fn_ast)\n",
    "        module = self.ib.get()\n",
    "        print(module)\n",
    "\n",
    "\n",
    "    def visit(self, node):\n",
    "        print(\"Visit \" + node.__class__.__name__)\n",
    "        return super().visit(node)\n",
    "    \n",
    "    def visit_Module(self, node: ast.Module):\n",
    "        if self.ir_module:\n",
    "            raise AssertionError(\"We should have only one module!\")\n",
    "        self.ir_module = I.ir_module()\n",
    "        with self.ir_module:\n",
    "            super().generic_visit(node)\n",
    "        \n",
    "    \n",
    "    def visit_FunctionDef(self, node: ast.FunctionDef):\n",
    "        fn = relax_builder.function()\n",
    "        self.entry = node.name\n",
    "        with fn:\n",
    "            R.func_name(node.name)\n",
    "            self._visit_compound_stmt(node.body)\n",
    "\n",
    "            if self.ret is None:\n",
    "                R.func_ret_value(rx.ShapeExpr([]))\n",
    "            else:\n",
    "                R.func_ret_value(self.ret)\n",
    "    \n",
    "    def visit_Pass(self, node: ast.Pass):\n",
    "        pass\n",
    "\n",
    "    def _visit_compound_stmt(self, stmts):\n",
    "        assert isinstance(stmts, (list, tuple))\n",
    "        for stmt in stmts:\n",
    "            ret = self.visit(stmt)\n",
    "            if ret is not None and isinstance(stmt, ast.Return):\n",
    "                self.ret = ret\n",
    "            \n",
    "    \n",
    "    def generic_visit(self, node: ast.AST):\n",
    "        raise NotImplementedError(\"Unsupported AST node type: {}\".format(type(node).__name__))\n",
    "\n",
    "\n",
    "@jit(target=\"cpu\")\n",
    "def add():\n",
    "    pass\n",
    "\n",
    "add()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "7ccfbc00",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-inf, inf)"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "float(\"-inf\")\n",
    "float.__class__.__name__\n",
    "# globals().values()\n",
    "float(\"-inf\"), float(\"+inf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "8cf0dfff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hello\n"
     ]
    }
   ],
   "source": [
    "for i in range(1):\n",
    "\tif True:\n",
    "\t\tpass\n",
    "\t\tprint(\"hello\")\n",
    "\t\tbreak\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "39cd7d83",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([2.2421, 1.5658, 1.5886, 0.0559])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "a = torch.randn(4)\n",
    "a * a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d4d88ca1",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "prod(): argument 'input' (position 1) must be Tensor, not list",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[10], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprod\u001b[49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mitem()\n",
      "\u001b[1;31mTypeError\u001b[0m: prod(): argument 'input' (position 1) must be Tensor, not list"
     ]
    }
   ],
   "source": [
    "torch.prod([1,2,3]).item()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my_env_torch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
